# Copyright 2025 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""expected output of ColumnParallelBatchedLinear"""

import numpy as np
import mindspore as ms
from mindformers.parallel_core.transformer_config import TransformerConfig


def get_data(transformer_config: TransformerConfig, seq_length=4, batch_size=1):
    """get state_dict and input_tensor for test"""
    np.random.seed(1)
    weight_dict = {
        "linear.weight": np.random.rand(transformer_config.num_moe_experts, transformer_config.ffn_hidden_size,
                                        transformer_config.hidden_size),
    }
    for k in weight_dict:
        weight_dict[k] = ms.Parameter(ms.tensor(weight_dict[k], dtype=ms.float32) / 100)

    input_numpy = np.random.rand(seq_length, batch_size, transformer_config.hidden_size)
    input_tmp = np.tile(input_numpy, (1, 2, 1))
    input_tensor = ms.tensor(input_tmp, dtype=ms.float32)

    return weight_dict, input_tensor


def get_output():
    """output for all cases"""
    gpu_data = np.array([[[0.0274658203, 0.0332031250, 0.0446777344, 0.0281982422,
                           0.0388183594, 0.0458984375, 0.0405273438, 0.0263671875,
                           0.0468750000, 0.0361328125, 0.0405273438, 0.0301513672,
                           0.0554199219, 0.0327148438, 0.0441894531, 0.0400390625]],
                         [[0.0292968750, 0.0478515625, 0.0539550781, 0.0273437500,
                           0.0510253906, 0.0534667969, 0.0485839844, 0.0454101562,
                           0.0371093750, 0.0373535156, 0.0375976562, 0.0383300781,
                           0.0517578125, 0.0324707031, 0.0537109375, 0.0541992188]],
                         [[0.0219726562, 0.0361328125, 0.0449218750, 0.0245361328,
                           0.0407714844, 0.0397949219, 0.0368652344, 0.0312500000,
                           0.0329589844, 0.0361328125, 0.0334472656, 0.0311279297,
                           0.0422363281, 0.0233154297, 0.0407714844, 0.0437011719]],
                         [[0.0314941406, 0.0405273438, 0.0520019531, 0.0301513672,
                           0.0473632812, 0.0529785156, 0.0444335938, 0.0417480469,
                           0.0383300781, 0.0434570312, 0.0429687500, 0.0324707031,
                           0.0544433594, 0.0291748047, 0.0529785156, 0.0541992188]],
                         [[0.0458984375, 0.0471191406, 0.0366210938, 0.0368652344,
                           0.0329589844, 0.0527343750, 0.0502929688, 0.0324707031,
                           0.0375976562, 0.0300292969, 0.0264892578, 0.0368652344,
                           0.0375976562, 0.0422363281, 0.0378417969, 0.0361328125]],
                         [[0.0505371094, 0.0529785156, 0.0439453125, 0.0419921875,
                           0.0385742188, 0.0588378906, 0.0541992188, 0.0380859375,
                           0.0388183594, 0.0432128906, 0.0349121094, 0.0522460938,
                           0.0480957031, 0.0505371094, 0.0415039062, 0.0393066406]],
                         [[0.0388183594, 0.0434570312, 0.0363769531, 0.0303955078,
                           0.0329589844, 0.0441894531, 0.0463867188, 0.0339355469,
                           0.0314941406, 0.0324707031, 0.0289306641, 0.0439453125,
                           0.0380859375, 0.0378417969, 0.0310058594, 0.0305175781]],
                         [[0.0495605469, 0.0505371094, 0.0466308594, 0.0446777344,
                           0.0395507812, 0.0610351562, 0.0605468750, 0.0393066406,
                           0.0463867188, 0.0427246094, 0.0346679688, 0.0510253906,
                           0.0424804688, 0.0456542969, 0.0449218750, 0.0400390625]]])
    golden_data = np.array([[[0.0274991784, 0.0332923606, 0.0447115004, 0.0281516593,
                              0.0387996919, 0.0458773635, 0.0404981486, 0.0263244994,
                              0.0468113869, 0.0361260101, 0.0405225083, 0.0301400106,
                              0.0555054806, 0.0328038260, 0.0442454368, 0.0401131921],
                             [0.0458219722, 0.0471253470, 0.0366109759, 0.0368310213,
                              0.0330453254, 0.0528750569, 0.0501868427, 0.0325499699,
                              0.0376858190, 0.0300634112, 0.0265092961, 0.0368209593,
                              0.0376662761, 0.0421527848, 0.0379402749, 0.0361749753]],
                            [[0.0292637348, 0.0477784052, 0.0539697073, 0.0273652002,
                              0.0509432107, 0.0535775386, 0.0485317595, 0.0455490164,
                              0.0371917523, 0.0373103581, 0.0376164801, 0.0383812636,
                              0.0518679284, 0.0325109288, 0.0536696576, 0.0541436188],
                             [0.0504607782, 0.0530052595, 0.0439509638, 0.0418924466,
                              0.0386122912, 0.0589769743, 0.0540959537, 0.0380703360,
                              0.0388395041, 0.0432056971, 0.0349408053, 0.0521817878,
                              0.0480804592, 0.0504504368, 0.0415362716, 0.0393412076]],
                            [[0.0219025128, 0.0361112952, 0.0448913909, 0.0244728178,
                              0.0407394320, 0.0398513563, 0.0369374864, 0.0311935302,
                              0.0328913815, 0.0362632424, 0.0335091762, 0.0311275702,
                              0.0421915799, 0.0233781300, 0.0408392958, 0.0436292104],
                             [0.0388612486, 0.0433428027, 0.0364659689, 0.0303557254,
                              0.0330081806, 0.0442704707, 0.0462704003, 0.0339764245,
                              0.0314693488, 0.0325502418, 0.0289373621, 0.0439760350,
                              0.0381023288, 0.0378445722, 0.0310597699, 0.0305675585]],
                            [[0.0314576961, 0.0405849814, 0.0520606041, 0.0301136170,
                              0.0474362746, 0.0529179946, 0.0444668345, 0.0418202356,
                              0.0384081490, 0.0436244272, 0.0429645404, 0.0324482918,
                              0.0545472354, 0.0290794689, 0.0530046560, 0.0541354828],
                             [0.0495874137, 0.0505451262, 0.0465402380, 0.0447414219,
                              0.0396549292, 0.0610013083, 0.0606010668, 0.0392509364,
                              0.0463327132, 0.0426730253, 0.0347297713, 0.0511451513,
                              0.0425956324, 0.0455601215, 0.0450493991, 0.0399953872]]])
    return gpu_data, golden_data
